import { Injectable, BadRequestException, NotFoundException, Logger } from '@nestjs/common';
import { InjectModel } from '@nestjs/mongoose';
import { Model, isValidObjectId } from 'mongoose';
import { Workflow } from './schemas/workflow.schema';
import { CreateWorkflowDto, UpdateWorkflowDto } from './dto/workflow.dto';
import { ComfyUIService } from '../comfyui/comfyui.service';

// 添加节点类型定义
interface WorkflowNode {
  class_type: string;
  inputs: {
    text?: string;
    seed?: number;
    steps?: number;
    cfg?: number;
    sampler_name?: string;
    width?: number;
    height?: number;
    batch_size?: number;
    [key: string]: any;  // 允许其他输入参数
  };
  [key: string]: any;  // 允许其他属性
}

interface WorkflowJson {
  [key: string]: WorkflowNode;
}

// 添加工作流格式类型定义
interface ComfyWorkflowNode {
  class_type: string;
  inputs: Record<string, any>;
  _meta?: {
    title?: string;
    description?: string;
    position?: { x: number; y: number };
  };
}

interface ComfyWorkflow {
  [key: string]: ComfyWorkflowNode;
}

@Injectable()
export class WorkflowService {
  private readonly logger = new Logger(WorkflowService.name);

  constructor(
    @InjectModel(Workflow.name) private workflowModel: Model<Workflow>,
    private readonly comfyuiService: ComfyUIService
  ) { }

  // 执行工作流
  async execute(id: string, params: Record<string, any>) {
    if (!id || !isValidObjectId(id)) {
      throw new NotFoundException('工作流不存在');
    }

    const workflow = await this.workflowModel.findById(id);
    if (!workflow) {
      throw new NotFoundException('工作流不存在');
    }

    return {
      workflow,
      params,
      status: 'success'
    };
  }

  // 生成图片
  async generateImage(workflowId: string, params: any) {
    try {
      // 验证工作流ID格式
      if (!isValidObjectId(workflowId)) {
        throw new BadRequestException('无效的工作流ID格式');
      }

      // 查找工作流
      const workflow = await this.workflowModel.findById(workflowId);
      if (!workflow) {
        throw new NotFoundException('工作流不存在');
      }

      if (!workflow.workflow_json) {
        throw new BadRequestException('工作流配置不存在');
      }

      this.logger.debug('Found workflow:', {
        id: workflowId,
        name: workflow.name,
        hasConfig: !!workflow.workflow_json
      });

      // 合并参数到工作流
      const mergedWorkflow = this.mergeParamsToWorkflow(workflow.workflow_json, params);

      //this.logger.debug('Submitting workflow to ComfyUI:', { mergedWorkflow });

      // 提交任务到 ComfyUI
      const result = await this.comfyuiService.executeWorkflow(mergedWorkflow);

      // 检查返回结果
      if (!result?.prompt_id) {
        throw new Error('提交任务失败: 未获取到 prompt_id');
      }

      // 检查任务状态
      const status = await this.comfyuiService.checkTaskStatus(result.prompt_id);

      // 返回任务信息给前端
      return {
        success: true,
        prompt_id: result.prompt_id,
        client_id: this.comfyuiService.client_id,
        message: '任务已提交，请等待生成结果',
        status: status.status || 'pending',
        exec_info: result.exec_info,
        node_errors: result.node_errors || [],
        node_warnings: result.node_warnings || [],
        queue_remaining: result.exec_info?.queue_remaining || 0,
        progress: status.progress || 0,
        outputs: status.outputs || null,
        current_node: status.node
      };

    } catch (error) {
      this.logger.error('生成图片失败:', error);
      return {
        success: false,
        message: error instanceof Error ? error.message : '未知错误',
        status: 'error',
        error: error instanceof Error ? error.message : '未知错误'
      };
    }
  }

  // 合并参数到工作流配置
  private mergeParamsToWorkflow(workflowJson: WorkflowJson, params: Record<string, any>) {
    try {
      this.logger.debug('Merging params into workflow:', { params });

      // 深拷贝工作流配置
      const workflow = JSON.parse(JSON.stringify(workflowJson)) as WorkflowJson;

      // 遍历所有节点
      for (const [nodeId, node] of Object.entries(workflow)) {
        if (!node || !node.class_type || !node.inputs) {
          continue;
        }

        // 根据节点类型处理参数
        switch (node.class_type) {
          case 'CLIPTextEncode':
            // 检查是正向还是负向提示词节点
            if (this.isPositivePromptNode(node, workflow)) {
              node.inputs.text = params.prompt || node.inputs.text;
              this.logger.debug(`Updated positive prompt for node ${nodeId}: ${node.inputs.text}`);
            } else if (this.isNegativePromptNode(node, workflow)) {
              node.inputs.text = params.negative_prompt || node.inputs.text;
              this.logger.debug(`Updated negative prompt for node ${nodeId}: ${node.inputs.text}`);
            }
            break;

          case 'EmptyLatentImage':
            // 处理图片尺寸
            if (params.size) {
              const [width, height] = params.size.split('x').map(Number);
              if (width && height) {
                node.inputs.width = width;
                node.inputs.height = height;
              }
            }
            node.inputs.batch_size = parseInt(params.batch_size) || 1;
            break;

          case 'KSampler':
            // 处理采样参数，添加类型转换和默认值
            node.inputs.seed = typeof params.seed === 'number' ? params.seed : -1;
            node.inputs.steps = params.steps ? parseInt(params.steps) : node.inputs.steps;
            node.inputs.cfg = params.cfg_scale ? parseFloat(params.cfg_scale) : node.inputs.cfg;
            node.inputs.sampler_name = params.sampler || node.inputs.sampler_name;
            break;
        }
      }

      //this.logger.debug('Merged workflow:', workflow);
      return workflow;
    } catch (error) {
      this.logger.error('Error merging params into workflow:', error);
      throw new Error('参数合并失败: ' + (error instanceof Error ? error.message : '未知错误'));
    }
  }

  // 添加辅助方法来判断是否为正向提示词节点
  private isPositivePromptNode(node: WorkflowNode, workflow: WorkflowJson): boolean {
    // 获取节点的ID
    const nodeId = Object.entries(workflow).find(([id, n]) => n === node)?.[0];
    if (!nodeId) return false;

    // 查找 KSampler 节点，检查其 positive 输入是否连接到当前节点
    return Object.values(workflow).some(otherNode =>
      otherNode.class_type === 'KSampler' &&
      otherNode.inputs?.positive?.[0] === nodeId
    );
  }

  // 添加辅助方法来判断是否为负向提示词节点
  private isNegativePromptNode(node: WorkflowNode, workflow: WorkflowJson): boolean {
    // 获取节点的ID
    const nodeId = Object.entries(workflow).find(([id, n]) => n === node)?.[0];
    if (!nodeId) return false;

    // 查找 KSampler 节点，检查其 negative 输入是否连接到当前节点
    return Object.values(workflow).some(otherNode =>
      otherNode.class_type === 'KSampler' &&
      otherNode.inputs?.negative?.[0] === nodeId
    );
  }

  // 创建工作流
  async create(createWorkflowDto: CreateWorkflowDto) {
    const workflow = new this.workflowModel(createWorkflowDto);
    return await workflow.save();
  }

  // 获取工作流列表
  async findAll() {
    return await this.workflowModel.find().sort({ created_at: -1 });
  }

  // 获取工作流详情
  async findOne(id: string) {
    if (!id || !isValidObjectId(id)) {
      throw new NotFoundException('工作流不存在');
    }

    const workflow = await this.workflowModel.findById(id);
    if (!workflow) {
      throw new NotFoundException('工作流不存在');
    }
    return workflow;
  }

  // 更新工作流
  async update(id: string, updateWorkflowDto: UpdateWorkflowDto) {
    if (!id || !isValidObjectId(id)) {
      throw new NotFoundException('工作流不存在');
    }

    try {
      // 获取原始工作流
      const originalWorkflow = await this.workflowModel.findById(id);
      if (!originalWorkflow) {
        throw new NotFoundException('工作流不存在');
      }

      // 如果提供了新的工作流配置，验证并保持格式
      let workflow_json = originalWorkflow.workflow_json;
      if (updateWorkflowDto.workflow_json) {
        workflow_json = this.validateAndFormatWorkflow(
          updateWorkflowDto.workflow_json,
          originalWorkflow.workflow_json
        );
      }

      const updatedWorkflow = await this.workflowModel.findByIdAndUpdate(
        id,
        {
          ...updateWorkflowDto,
          workflow_json,
          updated_at: new Date()
        },
        { new: true }
      );

      if (!updatedWorkflow) {
        throw new NotFoundException('工作流不存在');
      }

      return updatedWorkflow;
    } catch (error) {
      console.error('数据库更新失败:', {
        error,
        dto: updateWorkflowDto,
        stack: error.stack
      });
      throw error;
    }
  }

  // 验证并格式化工作流
  private validateAndFormatWorkflow(
    newWorkflow: any,
    originalWorkflow: ComfyWorkflow | null
  ): ComfyWorkflow {
    try {
      //console.log('原始工作流:', JSON.stringify(originalWorkflow, null, 2));
      // console.log('新工作流:', JSON.stringify(newWorkflow, null, 2));

      if (!newWorkflow || typeof newWorkflow !== 'object') {
        throw new Error('无效的工作流配置');
      }

      const formattedWorkflow: ComfyWorkflow = {};

      // 遍历新工作流的节点
      for (const [nodeId, node] of Object.entries(newWorkflow)) {
        if (!this.isValidNode(node)) {
          throw new Error(`节点 ${nodeId} 格式无效`);
        }

        // 获取原始节点的元数据
        const originalNode = originalWorkflow?.[nodeId];
        const originalMeta = originalNode?._meta;

        formattedWorkflow[nodeId] = {
          class_type: node.class_type,
          inputs: { ...node.inputs },
          _meta: {
            ...originalMeta, // 保持原有元数据
            ...(node._meta || {}) // 允许更新元数据
          }
        };
      }

      return formattedWorkflow;
    } catch (error) {
      this.logger.error('工作流验证失败:', {
        error: error.message,
        stack: error.stack,
        input: newWorkflow
      });
      throw error;
    }
  }

  // 验证节点格式
  private isValidNode(node: any): node is ComfyWorkflowNode {
    return (
      node &&
      typeof node === 'object' &&
      typeof node.class_type === 'string' &&
      node.inputs &&
      typeof node.inputs === 'object'
    );
  }

  // 删除工作流
  async remove(id: string) {
    if (!id || !isValidObjectId(id)) {
      throw new NotFoundException('工作流不存在');
    }

    const workflow = await this.workflowModel.findByIdAndDelete(id);
    if (!workflow) {
      throw new NotFoundException('工作流不存在');
    }
    return workflow;
  }

  // 添加检查生成状态的方法
  async checkGenerationStatus(promptId: string): Promise<any> {
    try {
      this.logger.debug('Checking generation status:', promptId);

      // 使用 ComfyUI 服务检查状态
      const status = await this.comfyuiService.checkTaskStatus(promptId);

      this.logger.debug('Generation status result:', status);

      // 如果任务完成，获取历史记录
      if (status.status === 'completed') {
        try {
          const history = await this.comfyuiService.getHistory(promptId);
          return {
            status: 'completed',
            outputs: history?.outputs || null,
            progress: 100
          };
        } catch (historyError) {
          this.logger.error('Failed to get history:', historyError);
          throw new Error('获取生成结果失败');
        }
      }

      // 如果任务失败
      if (status.status === 'failed') {
        return {
          status: 'failed',
          error: status.error || '生成失败',
          progress: 0
        };
      }

      // 如果任务仍在进行中
      return {
        status: status.status,
        progress: status.progress || 0,
        node: status.node
      };

    } catch (error) {
      this.logger.error('Check generation status failed:', {
        promptId,
        error: error instanceof Error ? error.message : '未知错误'
      });

      throw new Error(
        error instanceof Error
          ? error.message
          : '检查生成状态失败'
      );
    }
  }

  // 添加获取历史记录的辅助方法
  private async getGenerationHistory(promptId: string) {
    try {
      const history = await this.comfyuiService.getHistory(promptId);
      if (!history) {
        throw new Error('未找到生成历史记录');
      }
      return history;
    } catch (error) {
      this.logger.error('Get generation history failed:', {
        promptId,
        error: error instanceof Error ? error.message : '未知错误'
      });
      throw error;
    }
  }
} 